from langchain_community.vectorstores import Chroma
from langchain_core.documents import Document
from langchain_community.embeddings.dashscope import DashScopeEmbeddings
import json
import os

# 1. 加载JSON数据（使用您提供的知识库格式）
with open("data.json", "r", encoding="utf-8") as f:
    data = json.load(f)

# 2. 转换为Document对象列表
documents = []
for domain, qa_list in data.items():  # 遍历IT行业、医疗行业、人力资源
    for qa in qa_list:
        # 合并问题和答案作为文档内容
        content = f"问题：{qa['question']}\n答案：{qa['answer']}"
        # 创建Document对象，添加领域作为元数据
        documents.append(
            Document(
                page_content=content,
                metadata={
                    "domain": domain,
                    "question": qa["question"],  # 单独保留问题字段便于检索
                    "answer": qa["answer"]  # 单独保留答案字段便于检索
                }
            )
        )

# 3. 初始化向量数据库并添加数据
embeddings = DashScopeEmbeddings(
    model='text-embedding-v1',
    dashscope_api_key=os.getenv("DASHSCOPE_API_KEY")  # 请确保已设置环境变量
)

# 创建并持久化向量存储
vectordb = Chroma.from_documents(
    documents=documents,
    embedding=embeddings,
    persist_directory="./chroma_db"  # 向量数据库存储路径
)
